package com.lan.ai;

import com.alibaba.cloud.ai.autoconfigure.dashscope.*;
import com.alibaba.cloud.ai.dashscope.common.DashScopeApiConstants;
import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import com.lan.ai.mutimodelEmbedding.MutimodeDashScopeApi;
import io.micrometer.observation.ObservationRegistry;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.MultiValueMap;
import org.springframework.util.StringUtils;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * @author lanyanhua
 * @date 2025/1/21 15:22
 * @description
 */

@Slf4j
@Configuration
@EnableConfigurationProperties({ DashScopeConnectionProperties.class, DashScopeChatProperties.class,
        DashScopeImageProperties.class, DashScopeSpeechSynthesisProperties.class,
        DashScopeAudioTranscriptionProperties.class, DashScopeEmbeddingProperties.class,
        DashScopeRerankProperties.class })
public class DashScopeAutoConfig {


    @Bean
    public DashScopeEmbeddingModel dashScopeMultiModalEmbedding(DashScopeConnectionProperties commonProperties,
                                                           DashScopeEmbeddingProperties embeddingProperties, RestClient.Builder restClientBuilder,
                                                           WebClient.Builder webClientBuilder, RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler,
                                                           ObjectProvider<ObservationRegistry> observationRegistry,
                                                           ObjectProvider<EmbeddingModelObservationConvention> observationConvention) {

        var dashScopeApi = dashscopeEmbeddingApi(commonProperties, embeddingProperties, restClientBuilder,
                webClientBuilder, responseErrorHandler);

        var embeddingModel = new DashScopeEmbeddingModel(dashScopeApi, embeddingProperties.getMetadataMode(),
                embeddingProperties.getOptions(), retryTemplate,
                observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));

        observationConvention.ifAvailable(embeddingModel::setObservationConvention);

        return embeddingModel;
    }

    public MutimodeDashScopeApi dashscopeEmbeddingApi(DashScopeConnectionProperties commonProperties,
                                                      DashScopeEmbeddingProperties embeddingProperties, RestClient.Builder restClientBuilder,
                                                      WebClient.Builder webClientBuilder, ResponseErrorHandler responseErrorHandler) {
        ResolvedConnectionProperties resolved = resolveConnectionProperties(commonProperties, embeddingProperties,
                "embedding");
        //todo 替换默认的嵌入模型
        return new MutimodeDashScopeApi(resolved.baseUrl(), resolved.apiKey(), resolved.workspaceId(), restClientBuilder,
                webClientBuilder, responseErrorHandler);
    }

    private record ResolvedConnectionProperties(String baseUrl, String apiKey, String workspaceId,
                                                MultiValueMap<String, String> headers) {
    }

    private static @NotNull ResolvedConnectionProperties resolveConnectionProperties(
            DashScopeConnectionProperties commonProperties, DashScopeEmbeddingProperties modelProperties, String modelType) {

        String baseUrl = StringUtils.hasText(modelProperties.getBaseUrl()) ? modelProperties.getBaseUrl()
                : commonProperties.getBaseUrl();
        String apiKey = StringUtils.hasText(modelProperties.getApiKey()) ? modelProperties.getApiKey()
                : commonProperties.getApiKey();
        String workspaceId = StringUtils.hasText(modelProperties.getWorkspaceId()) ? modelProperties.getWorkspaceId()
                : commonProperties.getWorkspaceId();

        Map<String, List<String>> connectionHeaders = new HashMap<>();
        if (StringUtils.hasText(workspaceId)) {
            connectionHeaders.put("DashScope-Workspace", List.of(workspaceId));
        }

        // get apikey from system env.
        if (Objects.isNull(apiKey)) {
            if (Objects.nonNull(System.getenv(DashScopeApiConstants.DASHSCOPE_API_KEY))) {
                apiKey = System.getenv(DashScopeApiConstants.DASHSCOPE_API_KEY);
            }
        }

        Assert.hasText(baseUrl,
                "DashScope base URL must be set.  Use the connection property: spring.ai.dashscope.base-url or spring.ai.dashscope."
                        + modelType + ".base-url property.");
        Assert.hasText(apiKey,
                "DashScope API key must be set. Use the connection property: spring.ai.dashscope.api-key or spring.ai.dashscope."
                        + modelType + ".api-key property.");

        return new ResolvedConnectionProperties(baseUrl, apiKey, workspaceId,
                CollectionUtils.toMultiValueMap(connectionHeaders));
    }

}
